import re
from typing import Generator, List, Dict, Set, Iterator, Optional, Literal, AsyncIterator
from pydantic import BaseModel

from langchain_core.messages import BaseMessageChunk
from ..models.graph import GraphNodeParam, AgentStepLog, GraphState, Usage


class XmlToolResultParser:
    tool_thinking: str = "thinking"
    tools_names: List[str] = [tool_thinking]
    tools_params: Dict[str, List] = {}
    tools_stream: Dict[str, bool] = {}

    def __init__(self, tools: Dict[str, dict]):
        self.tools_names.extend(tools.keys())
        self.tools_params = {}
        self.tools_stream = {}
        for name, info in tools.items():
            schema = info.get("schema")
            stream = info.get("stream", True)
            self.tools_stream[name] = stream
            if isinstance(schema, type) and issubclass(schema, BaseModel):
                self.tools_params[name] = [field for field in schema.__fields__.keys()]
            else:
                self.tools_params[name] = schema

    async def parse_stream(self, chunks: AsyncIterator[BaseMessageChunk],
                           state: GraphState,
                           model: str = "default") -> (Generator)[
        tuple[AgentStepLog, str], None, None]:

        def build_agent_step_log(output, finish: bool = False,
                                 finish_reason: Optional[Literal["stop", "length"]] = "stop",
                                 usages: List[Usage] = []) -> AgentStepLog:
            _usages = list(map(lambda x: x.__dict__, usages))
            if isinstance(output, str):
                return AgentStepLog.build_answer(chunk=output, meta={"finish": finish,
                                                                     "finish_reason": finish_reason,
                                                                     "usages": _usages})
            return AgentStepLog.build_tool_action(action=output.type, output=output,
                                                  meta={"finish": finish,
                                                        "finish_reason": finish_reason,
                                                        "usages": _usages})

        last_line = ""
        processed: Set[str] = set()

        async for chunk in chunks:
            last_line += chunk.content

            if hasattr(chunk, 'usage_metadata') and chunk.usage_metadata:
                usage_meta = chunk.usage_metadata
                if 'input_tokens' in usage_meta and 'output_tokens' in usage_meta:
                    state.add_usage(
                        model=model,
                        input_tokens=usage_meta.get('input_tokens', 0),
                        output_tokens=usage_meta.get('output_tokens', 0)
                    )

            for tool_result, is_content_end in self.parse(content=last_line):
                if not tool_result:
                    continue

                if isinstance(tool_result, GraphNodeParam) and not self.tools_stream.get(tool_result.type, True):
                    if not is_content_end:
                        continue

                result_str = f"{tool_result}_{is_content_end}"
                if result_str in processed:
                    continue

                yield build_agent_step_log(tool_result, is_content_end, usages=state.usages), last_line
                processed.add(result_str)

        if last_line:
            tool_names = self.content_contain_tools(last_line)
            if not tool_names:
                if last_line not in processed:
                    yield build_agent_step_log(last_line, True, usages=state.usages), last_line
                    processed.add(last_line)
            else:
                last_end_idx = max(last_line.rfind(f"</{name}>") for name in tool_names)
                if last_end_idx >= 0:
                    remaining = last_line[last_end_idx + last_line[last_end_idx:].find('>') + 1:].strip()
                    if remaining and remaining not in processed:
                        ## 处理超长中断的情况
                        remaining_tool_names = self.content_contain_tools(remaining)
                        if remaining_tool_names:
                            # 如果剩余内容包含工具，说明是超长导致的中断
                            tool_result, _ = next(self.parse(remaining))
                            if tool_result and (not isinstance(tool_result, GraphNodeParam) or tool_result.params):
                                yield build_agent_step_log(tool_result, True, "length", usages=state.usages), last_line
                        else:
                            yield build_agent_step_log(remaining, True, usages=state.usages), last_line

    def parse(self, content: str, buffer_size: int = 30) -> Generator[
        tuple[GraphNodeParam | str | None, bool], None, None]:
        if len(content) < buffer_size:
            yield None, False
            return

        tool_names = self.content_contain_tools(content)
        if not tool_names:
            yield content[:len(content) - buffer_size], False
            return

        remaining_content = content
        last_content_end = False
        for tool_name in tool_names:
            for tool_result, content_end in self._parse_content(tool_name, remaining_content):
                if not tool_result or (isinstance(tool_result, GraphNodeParam) and not tool_result.params):
                    continue
                yield tool_result, content_end
                last_content_end = content_end
                if content_end:
                    symbol_end = f"</{tool_name}>"
                    end_idx = remaining_content.find(symbol_end)
                    if end_idx != -1:
                        remaining_content = remaining_content[end_idx + len(symbol_end):].strip()

        if last_content_end and len(remaining_content) > buffer_size:
            tool_names = self.content_contain_tools(remaining_content)
            if not tool_names:
                yield remaining_content[:len(remaining_content) - buffer_size], False

    def _parse_content(self, tool_name: str, content: str) -> Generator[
        tuple[GraphNodeParam | str | None, bool], None, None]:
        symbol_start = f"<{tool_name}>"
        symbol_end = f"</{tool_name}>"
        if symbol_start not in content or len(content) < len(symbol_start) + len(symbol_end):
            yield None, False
            return

        params = []
        if tool_name != "thinking":
            if tool_name in self.tools_params:
                params = self.tools_params[tool_name]

        start_idx = content.find(symbol_start)
        if start_idx > 0:
            prefix_content = content[:start_idx].strip()
            if prefix_content:
                yield prefix_content, True
        content = content[start_idx + len(symbol_start):]

        if symbol_end in content:
            content_end = True
            content = content[:content.find(symbol_end)]
        else:
            content_end = False
            content = content[:-len(symbol_end)]

        if not params:
            yield content, content_end
            return

        tool_result = GraphNodeParam(type=tool_name, params={})

        for param in params:
            param_start = f"<{param}>"
            param_end = f"</{param}>"
            if param_start not in content:
                continue

            start_idx = content.find(param_start)
            end_idx = content.find(param_end, start_idx)

            if end_idx == -1:
                param_content = content[start_idx + len(param_start):-len(param_end)].strip()
            else:
                param_content = content[start_idx + len(param_start):end_idx].strip()

            if param_content:
                tool_result.add_param(param, param_content)

        yield tool_result, content_end

    def content_contain_tools(self, content: str, direct_break: bool = False, except_thinking: bool = False) -> List[
        str]:
        tool_names = []
        tags = re.findall(r'<([^/][^>]*)>', content)
        for tag in tags:
            if tag in self.tools_names:
                if except_thinking and tag == self.tool_thinking:
                    continue
                tool_names.append(tag)
                if direct_break:
                    break
        return tool_names
